from abc import ABC, abstractmethod
import os, io, time
import logging
import wave
import numpy as np
from datetime import timedelta
from scipy.io.wavfile import write
import shutil
from services.utils import base64_str_encode, base64_str_decode
import asyncio
from aiocache import Cache
from aioredis import from_url, Redis
from typing import Union
try:
    from minio import Minio
    from minio.helpers import ObjectWriteResult
    from minio.datatypes import Object as ReadObject
    from minio.error import S3Error
except ModuleNotFoundError:
    pass

logger = logging.getLogger(__name__)


class BaseStorage(ABC):
    """
    抽像的存储
    Args:
        ABC (_type_): _description_
    """

    # 存储
    @abstractmethod
    def upload(self, directory: str, reader: io.BufferedReader | np.ndarray, extension: str, **kwargs,) -> tuple[str, int]:
        """
        小文件上传. kwargs在extension=wav时支持: rate/音频采样率, channel/通道数
        Args:
            directory (str): 保存的子目录名
            reader (io.BufferedReader | np.ndarray): 上传的数据
            extension (str): 保存文件的扩展名

        Returns:
            tuple[str,int]: [0]文件的相对路径, [1]文件的大小. 上传失败: [0]=None, [1]=0
        """
        pass

    # 位置/可能是本地全路径/可能是网络路径
    @abstractmethod
    def getpath(self, relative_path: str) -> str | None:
        """
        返回保存文件的可访问路径
        Args:
            relative_path (str): 文件的相对路径

        Returns:
            str|None: 可访问文件路径
        """
        pass

    @abstractmethod
    def geturl(self, relative_path: str) -> str | None:
        """
        返回保存文件的可访问url
        Args:
            relative_path (str): 文件的相对路径

        Returns:
            str|None: 文件的可访问url
        """
        pass

    # getpath的复数形式
    @abstractmethod
    def listpath(self, relative_path_list: list[str], save_local: bool = True) -> list[str]:
        """
        getpath的复数形式
        Args:
            relative_path_list (list[str]): 文件的相对路径列表
            save_local (bool, optional): 是否保存到本地. Defaults to True.

        Returns:
            list[str]: 可访问路径列表
        """
        pass

    # 是否存在
    @abstractmethod
    def exists(self, relative_path: str) -> bool:
        """
        是否存在
        Args:
            relative_path (str): 文件的相对路径

        Returns:
            bool: 存在返回True
        """
        pass

    # 移除
    @abstractmethod
    def remove(self, relative_path: str) -> None:
        """
        移除指定文件
        Args:
            relative_path (str): 文件的相对路径
        """
        pass

    # 流式
    @abstractmethod
    def stream(self, directory: str, file_path: str) -> str:
        """
        流式大文件上传
        Args:
            directory (str): 保存的子目录名
            file_path (str): 待上传文件的路径

        Returns:
            str: 文件的相对路径/
        """
        pass

    @staticmethod
    def generate_file_name(extension: str) -> str:
        return "{}.{}".format(str(int(time.time())), extension)

    @staticmethod
    def generate_relative_path(filename: str, directory: str) -> str:
        return "{}/{}".format(directory, filename)

    @staticmethod
    def build_relative_path(directory: str, extension: str) -> str:
        return BaseStorage.generate_relative_path(
            BaseStorage.generate_file_name(extension), directory
        )

    @staticmethod
    def relative_path_template(directory: str, filename: str, extension: str) -> str:
        return "{}/{}.{}".format(directory, filename, extension)


class LocalStorageURLCoder(ABC):
    @abstractmethod
    def encode(self, relative_path: str) -> str:
        """
        路径编码
        Args:
            relative_path (str): 文件的相对路径

        Returns:
            str: 生成文件路径散列码
        """
        pass

    @abstractmethod
    def decode(self, path_hashcode: str) -> str | None:
        """
        路径解码
        Args:
            path_hashcode (str): 文件路径散列码

        Returns:
            str|None: 文件可访问的本地路径
        """
        pass


class LocalFileStorage(BaseStorage, LocalStorageURLCoder):
    """
    使用站内文件存储
    Args:
        BaseStorage (_type_): _description_
    """

    def __init__(self, base_path: str = ".", base_url: str = "/oss"):
        """
        初始化
        Args:
            base_path (str, optional): 存储的基路径. Defaults to '.'.
            base_url (str, optional): url的前缀部分. Defaults to None.
        """
        self.basePath = base_path
        self.baseUrl = base_url

    def upload(self, directory: str, reader: io.BufferedReader | np.ndarray, extension: str, **kwargs,) -> tuple[str, int]:
        if extension and extension.lower() == "wav":
            return self.__write_wav_file(directory, reader, **kwargs)

        size: int = 0
        saveDirectoryPath = LocalFileStorage.build_relative_path(directory, extension)
        with open(os.path.join(self.basePath, saveDirectoryPath), "wb") as f:
            while True:
                chunk = reader.read(8192)  # 8192 bytes by default.
                if not chunk:
                    break
                size += len(chunk)
                f.write(chunk)

        return saveDirectoryPath, size

    def __write_wav_file(self, directory: str, reader: io.BufferedReader | np.ndarray, **kwargs) -> tuple[str, int]:
        """
        上传wav文件
        Args:
            directory (str): 保存的子目录名
            reader (io.BufferedReader | np.ndarray): 上传的数据

        Returns:
            tuple[str,int]: [0]文件的相对路径, [1]文件的大小; 若非允许的类型(io.BufferedReader|np.ndarray)返回: [0]None, [1]0
        """
        saveDirectoryPath = LocalFileStorage.build_relative_path(directory, "wav")
        sample_rate: int = kwargs.pop("rate", 48000)
        channel_count: int = kwargs.pop("channel", 2)
        # @see: https://blog.finxter.com/converting-python-numpy-arrays-to-wav-files-top-solutions/
        allowed: bool = True
        # Open a WAV file
        save_wav_path = os.path.join(self.basePath, saveDirectoryPath)
        with wave.open(save_wav_path, "w") as wav_file:
            # Define audio parameters
            wav_file.setnchannels(channel_count)  # Mono
            wav_file.setsampwidth(2)  # Two bytes per sample
            wav_file.setframerate(sample_rate)
            if isinstance(reader, np.ndarray):
                wav_file.writeframes(reader.tobytes())
            elif isinstance(reader, io.BufferedReader):
                # Convert the NumPy array to bytes and write it to the WAV file
                wav_file.writeframesraw(reader.raw.getvalue())
            else:
                allowed = False

        if allowed:
            return saveDirectoryPath, os.path.getsize(save_wav_path)
        else:
            return None, 0

    def getpath(self, relative_path: str) -> str | None:
        return (self.__build_file_path(relative_path) if self.exists(relative_path) else None)

    def geturl(self, relative_path: str) -> str | None:
        path_hashcode = self.encode(relative_path)
        if path_hashcode is None or self.baseUrl is None:
            return None

        return "{}/{}".format(self.baseUrl, path_hashcode)

    def encode(self, relative_path: str) -> str:
        return base64_str_encode(relative_path)

    def decode(self, path_hashcode: str) -> str | None:
        relative_path = base64_str_decode(path_hashcode)
        return self.getpath(relative_path)

    def listpath(self, relative_path_list: list[str], save_local: bool = True) -> list[str]:
        return [ fp for fp in [self.getpath(rp, save_local) for rp in relative_path_list] if fp is not None]

    def exists(self, relative_path: str) -> bool:
        return os.path.exists(self.__build_file_path(relative_path))

    def remove(self, relative_path: str) -> None:
        fp = self.getpath(relative_path)
        if fp is not None:
            return os.remove(fp)

    def stream(self, directory: str, file_path: str) -> str:
        _, file_extension = os.path.splitext(file_path)
        saveDirectoryPath = LocalFileStorage.build_relative_path(directory, file_extension[1:])
        #
        shutil.copy2(file_path, self.__build_file_path(saveDirectoryPath))
        return saveDirectoryPath

    def __build_file_path(self, relative_path: str) -> str:
        return "{}/{}".format(self.basePath, relative_path)


class MinIoObjectStorage(BaseStorage):
    """
    使用MinIO对象存储
    Args:
        BaseStorage (_type_): _description_
    """

    def __init__(self, endpoint: str, access_key: str, secret_key: str, is_https: bool = False, bucket: str = "ahhs",):
        self.__client = Minio(endpoint=endpoint, access_key=access_key, secret_key=secret_key, secure=is_https,)
        self.__exists_bucket(bucket)
        self.bucket = bucket

    def upload(self, directory: str, reader: io.BufferedReader | np.ndarray, extension: str, **kwargs,) -> tuple[str, int]:
        if extension and extension.lower() == "wav":
            return self.__write_wav_file(directory, reader, **kwargs)

        saveDirectoryPath = MinIoObjectStorage.build_relative_path(directory, extension)
        #
        assert isinstance(reader, io.BufferedReader), "Numpy array only support wav extension"
        object_bytes, object_data = self.__calc_buffer_byte_size(reader)
        # ObjectWriteResult
        result: ObjectWriteResult = self.__client.put_object(
            bucket_name=self.bucket,
            object_name=saveDirectoryPath,
            data=object_data,
            length=object_bytes,
        )
        self.__log_debug("upload binary result:%s, object version:%s, etag:%s, location:%s", type(result), result.version_id, result.etag, result.location,)
        return saveDirectoryPath, object_bytes

    def getpath(self, relative_path: str) -> str | None:
        try:
            local_file_path = "./{}".format(relative_path)
            result: ReadObject = self.__client.fget_object(
                bucket_name=self.bucket,
                object_name=relative_path,
                file_path=local_file_path,
            )
            self.__log_debug("down object result:%s, object version:%s, etag:%s, size:%s", type(result), result.version_id, result.etag, result.size,)
            return local_file_path
        except S3Error as err:
            self.__log_debug("getpath has exception:%s, message:%s", repr(err), str(err))
            return None

    def geturl(self, relative_path: str) -> str | None:
        try:
            url = self.__client.presigned_get_object(
                bucket_name=self.bucket,
                object_name=relative_path,
                expires=timedelta(hours=2),
            )
            return url
        except S3Error as err:
            self.__log_debug("getpath has exception:%s, message:%s", repr(err), str(err))
            return None

    def listpath(self, relative_path_list: list[str], save_local: bool = True) -> list[str]:
        return list(filter(lambda fp: fp is not None, [self.getpath(rp, save_local) for rp in relative_path_list],))

    def exists(self, relative_path: str) -> bool:
        result = None
        try:
            result: ReadObject = self.__client.stat_object( bucket_name=self.bucket, object_name=relative_path)
            self.__log_debug("exists object result:%s, object version:%s, etag:%s, size:%s", type(result), result.version_id, result.etag, result.size,)
        except S3Error as err:
            self.__log_debug("exists object exception:%s, message:%s", repr(err), str(err))

        return result is not None

    def remove(self, relative_path: str) -> None:
        self.__client.remove_object(bucket_name=self.bucket, object_name=relative_path)

    def stream(self, directory: str, file_path: str) -> str:
        _, file_extension = os.path.splitext(file_path)
        saveDirectoryPath = MinIoObjectStorage.build_relative_path(directory, file_extension[1:])

        result: ObjectWriteResult = self.__client.put_object(
            bucket_name=self.bucket,
            object_name=saveDirectoryPath,
            data=io.FileIO(file_path),
            length=-1,
            part_size=1024 * 1024 * 5,
        )
        self.__log_debug("upload stream result:%s, object version:%s, etag:%s, location:%s", type(result), result.version_id, result.etag, result.location,)
        return saveDirectoryPath

    def __exists_bucket(self, bucket: str = "ahhs") -> None:
        assert self.__client.bucket_exists(bucket), "not find bucket: {}".format(bucket)

    def __write_wav_file(self, directory: str, reader: io.BufferedReader | np.ndarray, **kwargs) -> tuple[str, int]:
        saveDirectoryPath = MinIoObjectStorage.build_relative_path(directory, "wav")
        sample_rate: int = kwargs.pop("rate", 48000)
        # channel_count:int=kwargs.pop('channel', 2)

        # @see: https://stackoverflow.com/questions/51300281/how-to-convert-numpy-array-to-bytes-object-without-save-audio-file-on-disk
        object_output: io.BytesIO = None
        if isinstance(reader, np.ndarray):
            self.__log_debug("upload wav data source: numpy")
            bytes_wav = bytes()
            object_output = io.BytesIO(bytes_wav)
            write(object_output, sample_rate, reader)
        elif isinstance(reader, io.BufferedReader):
            self.__log_debug("upload wav data source: BufferedReader")
            # Convert the NumPy array to bytes and write it to the WAV file
            object_output = io.BytesIO(reader)
            write(object_output, sample_rate, reader)
        else:
            self.__log_debug("upload wav data source: Not Allow!!")
            return None, 0

        object_bytes = object_output.getbuffer().nbytes
        # ObjectWriteResult
        result: ObjectWriteResult = self.__client.put_object(
            bucket_name=self.bucket,
            object_name=saveDirectoryPath,
            data=object_output,
            length=object_bytes,
        )
        self.__log_debug("upload wav result:%s, object version:%s, etag:%s, location:%s", type(result), result.version_id, result.etag, result.location,)
        return saveDirectoryPath, object_bytes

    def __calc_buffer_byte_size(self, reader: io.BufferedReader) -> tuple[int, io.BytesIO]:
        size = 0
        bio = io.BytesIO()
        while True:
            chunk = reader.read(4096)
            if not chunk:
                break
            size += len(chunk)
            bio.write(chunk)

        bio.seek(0)
        self.__log_debug("buffer reader byte length:%d", size)
        return size, bio

    def __log_debug(self, msg: str, *args) -> None:
        logger.debug(f"[MinIoObjectStorage]{msg}", *args)


class AiocacheMemory:
    """
    内存存储
    """

    def __init__(self,):
        """
        初始化
        """
        self.storage = Cache(Cache.MEMORY)

    # 缓存指定key与其关联的字典
    async def set_key(self, key: str, value: any, tts_seconds:Union[int, None]=86400):
        """
        使用指定的key存储字典值
        Args:
            key (str): _description_
            value (dict): _description_
            tts_seconds (int, optional): 有效时长, 单位秒. Defaults to 86400.

        Returns:
            _type_: _description_
        """
        rs = await self.storage.set(key, value, ttl=tts_seconds)
        return rs

    # 是否存在指定的key
    async def exist_key(self, key: str) -> bool:
        """
        是否存在指定的key
        Args:
            key (str): _description_

        Returns:
            bool: 存在返回True
        """
        try:
            return await self.storage.exists(key)
        except asyncio.TimeoutError as e:
            return False

    # 删除指定key. 并返回关联的字典
    async def poll_key(self, key: str, defVal: any = None) -> any:
        """
        删除指定key. 并返回关联的字典
        Args:
            key (str): _description_
            defVal (dict, optional): 若key不可用时的默认字典. Defaults to {}.

        Returns:
            dict: _description_
        """
        ele = await self.get_key(key, defVal)
        await self.storage.delete(key)
        return ele

    # 返回指定key关联的字典
    async def get_key(self, key: str, defVal: any = None) -> any:
        """
        返回指定key关联的字典
        Args:
            key (str): _description_
            defVal (dict, optional): 若key不可用时的默认字典. Defaults to {}.

        Returns:
            dict: _description_
        """
        return await self.storage.get(key, default=defVal)
    
    async def del_key(self, key:str)->None:
        return await self.storage.delete(key)
    
class AioRedisCache:
    # @see: aioredis.connection.ConnectKwargs
    # redis_config 可选字典Key:
    # username
    # password
    # port
    # db
    def __init__(self, redis_url: str, **redis_config):
        self.__client_url = redis_url
        self.__client_config = redis_config
        self.__client:Redis = None
        
    async def _connect(self):
        self.__client = from_url(f"redis://{self.__client_url}", **self.__client_config)
    
    async def _disconnect(self, exc_type, exc_value, exc_tb):
        await self.__client.close()

    async def __aenter__(self):
        await self._connect()  # added
        return self
    
    async def __aexit__(self, exc_type, exc, tb):
        await self._disconnect(exc_type, exc, tb)
        
    async def set_key(self, key: str, value: str, tts_seconds:Union[int, None]=86400):
        await self.__client.set(name=key, value=value, ex=tts_seconds, nx=True)
        
    async def exist_key(self, key: str) -> bool:
        hasize:int = await self.__client.exists(key)
        return True if hasize > 0 else False

    async def poll_key(self, key: str, defVal: str = None) -> str:
        data = await self.__client.execute_command("getdel", key)
        return data if data else defVal
    
    async def get_key(self, key: str, defVal: str = None) -> str:
        data = await self.__client.get(key)
        return data if data else defVal
    
    async def del_key(self, key:str)->None:
        await self.__client.delete(key)

    async def add_set_key(self, key:str, *values:str):
        await self.__client.sadd(key, *values)
        
    async def exist_set_member(self, key:str, value:str)->bool:
        await self.__client.sismember(name=key, value=value)
        
    async def close(self):
        if self.__client:
            await self.__client.close()